"""
通过改写来写作
"""
import os
import random
import re
import time
import traceback
from typing import Dict

from langchain import PromptTemplate

try:
    from bot.openai_bot import OpenAIBot
    from conf.config import logger, BASE_DIR, config
    from db.building_dao import BuildingDAO
    from prompts.prompt_templates import BUILDING_INFO_BASE_PROMPT, ARTICLE_REWRITING_PROMPT
    from utils.constants import EN_TO_CN, SALE_STATUS
    from utils.format_article import count_words, all_has_tags, extract_abstract_and_keywords_and_content
    from utils.laod_data import load_excel_sheet, load_title_list_within_contexts
    from utils.output_file import output_txt, output_json
except ModuleNotFoundError:
    import os
    import sys
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))  # 离开IDE也能正常导入自己定义的包
    from bot.openai_bot import OpenAIBot
    from conf.config import logger, BASE_DIR, config
    from db.building_dao import BuildingDAO
    from prompts.prompt_templates import BUILDING_INFO_BASE_PROMPT, ARTICLE_REWRITING_PROMPT
    from utils.constants import EN_TO_CN, SALE_STATUS
    from utils.format_article import count_words, all_has_tags, extract_abstract_and_keywords_and_content
    from utils.laod_data import load_excel_sheet, load_title_list_within_contexts
    from utils.output_file import output_txt, output_json


def generate_article(prompt, filename, title):
    """
    生产文章
    :param prompt: 提交gpt的prompt
    :param filename: 保存的文件名
    :param title: 标题
    :return:
    """
    openai_bot = OpenAIBot()

    answer = openai_bot.reply(question=prompt)

    data = extract_abstract_and_keywords_and_content(answer)
    data.update({"id": str(filename)})
    data.update({"title": title})

    # 判断摘要字数是否在限值内
    abstract_word_count = count_words(data.get("abstract"), excluded_char_ist=[])
    if abstract_word_count > config.get("word_limit_for_abstract"):
        raise ValueError(f"摘要字数超限值，共{abstract_word_count}字")

    # 判断正文是否都有标签
    if not all_has_tags(text=data.get("content"), startswith_tags=config.get("startswith_tags"),
                        endswith_tags=config.get("endswith_tags")):
        raise ValueError(f"正文有些段落标签错误")

    # 输出到json(中文为不可读的Unicode编码)
    output_json(
        output_full_path=os.path.join(BASE_DIR, f"output/{filename}.json"),
        data=data, ensure_ascii=True
    )

    # prompt输出到txt文件
    output_txt(
        output_full_path=os.path.join(BASE_DIR, f'output/{filename}_prompt.txt'),
        txt=prompt
    )

    # 文章结果输出到txt文件
    output_txt(
        output_full_path=os.path.join(BASE_DIR, f"output/{filename}.txt"),
        txt=answer
    )


def main():
    fail_list = list()
    title_list = load_title_list_within_contexts()
    for item in title_list:
        filename = item.get("filename")
        title = item.get("title")
        contexts = item.get("contexts")
        try:
            # 渲染模板
            prompt = PromptTemplate(
                input_variables=["title", "contexts"],
                template=ARTICLE_REWRITING_PROMPT,
                template_format="jinja2"
            )
            prompt_formatted = prompt.format(title=title, contexts=contexts)

            generate_article(prompt=prompt_formatted, filename=f'{filename}', title=title)
            time.sleep(random.uniform(10, 15))

        except Exception:
            error_str = traceback.format_exc()
            logger.error(error_str)
            fail_list.append({"filename": filename, "log": error_str})

    logger.warning(f'失败列表：\n{[item.get("filename") for item in fail_list]}')
    logger.warning(f'失败列表详情：\n{fail_list}')


if __name__ == '__main__':
    main()
